package main

import (
	"bytes"
	"fmt"
	"io"
	"log"
	"os"
	"sync/atomic"
	"time"

	"gitee.com/general252/ebml-go"
	"gitee.com/general252/ebml-go/core/codecs"
	"gitee.com/general252/ebml-go/core/formats"
	"gitee.com/general252/ebml-go/matroska"
	"github.com/bluenviron/gortsplib/v3"
	"github.com/bluenviron/gortsplib/v3/pkg/formats/rtph264"
	"github.com/bluenviron/gortsplib/v3/pkg/formats/rtpmpeg4audiogeneric"
	"github.com/bluenviron/gortsplib/v3/pkg/media"
	"github.com/bluenviron/mediacommon/pkg/codecs/h264"
	"github.com/bluenviron/mediacommon/pkg/codecs/h265"
	"github.com/bluenviron/mediacommon/pkg/codecs/mpeg4audio"
	"github.com/deepch/vdk/av"
	"github.com/deepch/vdk/codec/aacparser"
	"github.com/deepch/vdk/codec/h264parser"
	"github.com/deepch/vdk/codec/h265parser"

	mp4gomedia "gitee.com/general252/gomedia/go-mp4"
	formatsvdk "github.com/bluenviron/gortsplib/v3/pkg/formats"
	mp4vdk "github.com/deepch/vdk/format/mp4"
)

func init() {
	log.SetFlags(log.LstdFlags | log.Lshortfile)
}

func main() {
	//testAVCC()
	//testAnnexB()
	ebmlDemo()
}

func testAVCC() {
	inFilename := `..\container\res\input.h264`
	fpInput, _ := os.Open(inFilename)
	defer fpInput.Close()

	fpOut, _ := os.Create("out-1.h264")
	defer fpOut.Close()

	err := codecs.ConvertAnnexBToAVCC(fpInput, fpOut)
	log.Println(err)
}

func testAnnexB() {
	inFilename := "out-1.h264"
	fpInput, _ := os.Open(inFilename)
	defer fpInput.Close()

	fpOut, _ := os.Create("out-2.h264")
	defer fpOut.Close()

	err := codecs.ConvertAVCCToAnnexB(fpInput, fpOut)
	log.Println(err)
}

func testMP41() {
	f, _ := os.Create("out1.mp4")
	defer f.Close()

	m, err := formats.MP4GoMediaMuxer(f)
	if err != nil {
		log.Println(err)
		return
	} else {
		defer func() {
			if err := m.WriteTrailer(); err != nil {
				log.Println(err)
			}
		}()
	}

	trackId := m.AddVideoTrack(mp4gomedia.MP4_CODEC_H264)

	inFilename := `..\container\res\input.h264`
	fp, err := os.Open(inFilename)
	if err != nil {
		log.Println(err)
		return
	} else {
		defer fp.Close()
	}

	var (
		timestamp time.Duration
		startCode = []byte{0x0, 0x0, 0x0, 0x01}
	)

	codecs.EmitNALUH264Reader(fp, codecs.NALUFormatAnnexB, codecs.NALUFormatNo, func(t h264.NALUType, data []byte) {
		switch t {
		case h264.NALUTypeSPS, h264.NALUTypePPS:
			// 7
			buff := append(startCode, data...)

			pts := uint64(timestamp.Milliseconds())
			if err = m.Write(trackId, buff, pts, pts); err != nil {
				log.Println(err)
				break
			}
		case h264.NALUTypeIDR, h264.NALUTypeNonIDR:
			// 1
			buff := append(startCode, data...)

			pts := uint64(timestamp.Milliseconds())
			if err = m.Write(trackId, buff, pts, pts); err != nil {
				log.Println(err)
				break
			}

			timestamp += 40 * time.Millisecond
		default:
		}
	})

}

func testMP42() {
	f, _ := os.Create("out2.mp4")
	defer f.Close()

	m, err := formats.MP4GoMediaMuxer(f)
	if err != nil {
		log.Println(err)
		return
	} else {
		defer func() {
			if err := m.WriteTrailer(); err != nil {
				log.Println(err)
			}
		}()
	}

	trackId := m.AddAudioTrack(mp4gomedia.MP4_CODEC_G711A, mp4gomedia.WithAudioChannelCount(1), mp4gomedia.WithAudioSampleRate(8000))

	inFilename := `..\container\res\output.g711a`
	fp, err := os.Open(inFilename)
	if err != nil {
		log.Println(err)
		return
	} else {
		defer fp.Close()
	}

	var timestamp time.Duration

	for {
		buff := make([]byte, 320)
		n, err := fp.Read(buff)
		if err != nil {
			break
		}

		pts := uint64(timestamp.Milliseconds())

		if err = m.Write(trackId, buff[:n], pts, pts); err != nil {
			log.Println(err)
			break
		}

		timestamp += 40 * time.Millisecond
	}
}

func testMP43() {
	f, _ := os.Create("out3.mp4")
	defer f.Close()

	m, err := formats.MP4GoMediaMuxer(f)
	if err != nil {
		log.Println(err)
		return
	} else {
		defer func() {
			if err := m.WriteTrailer(); err != nil {
				log.Println(err)
			}
		}()
	}

	trackId := m.AddAudioTrack(mp4gomedia.MP4_CODEC_AAC, mp4gomedia.WithAudioChannelCount(1), mp4gomedia.WithAudioSampleRate(8000))

	/////////////////////////
	var inFilename = `..\container\res\output.aac`
	var fp io.ReadCloser
	if fp, err = os.Open(inFilename); err != nil {
		log.Println(err)
		return
	} else {
		defer fp.Close()
	}

	var (
		samples = 0
	)

	_ = formats.EmitADTSReader(fp, func(packet *mpeg4audio.ADTSPacket, data []byte, head []byte) {
		timestamp := time.Duration(samples*1000/packet.SampleRate) * time.Millisecond
		log.Println(timestamp)

		pts := uint64(timestamp.Milliseconds())
		samples += 1024

		if err = m.Write(trackId, append(head, data...), pts, pts); err != nil {
			log.Println(err)
		}
	})

}

func testMP44() {
	f, _ := os.Create("out4.mp4")
	defer f.Close()

	m, err := formats.MP4GoMediaMuxer(f)
	if err != nil {
		log.Println(err)
		return
	} else {
		defer func() {
			if err := m.WriteTrailer(); err != nil {
				log.Println(err)
			} else {
				log.Println("close success")
			}
		}()
	}

	trackId := m.AddVideoTrack(mp4gomedia.MP4_CODEC_H264)
	var (
		writeVideo atomic.Bool
		writeAudio atomic.Bool
		isRun      = true
		topPTS     uint64
	)
	writeAudio.Store(true)

	go func() {

		inFilename := `..\container\res\input.h264`
		fp, err := os.Open(inFilename)
		if err != nil {
			log.Println(err)
			return
		} else {
			defer fp.Close()
		}

		var (
			timestamp time.Duration
			startCode = []byte{0x0, 0x0, 0x0, 0x01}
		)

		codecs.EmitNALUH264Reader(fp, codecs.NALUFormatAnnexB, codecs.NALUFormatNo, func(t h264.NALUType, data []byte) {
			pts := uint64(timestamp.Milliseconds())

			for isRun {
				if writeVideo.Load() {
					if pts <= topPTS {
						break
					} else {
						writeVideo.Store(false)
						writeAudio.Store(true)
					}
				}

				time.Sleep(time.Millisecond)
			}

			log.Println("video", timestamp)
			time.Sleep(time.Millisecond * 30)

			switch t {
			case h264.NALUTypeSPS, h264.NALUTypePPS:
				// 7
				buff := append(startCode, data...)

				if err = m.Write(trackId, buff, pts, pts); err != nil {
					log.Println(err)
					break
				}
			case h264.NALUTypeIDR, h264.NALUTypeNonIDR:
				// 1
				buff := append(startCode, data...)

				if err = m.Write(trackId, buff, pts, pts); err != nil {
					log.Println(err)
					break
				}

				timestamp += 40 * time.Millisecond
			default:
			}
		})

		isRun = false

	}()

	{
		trackIdAudio := m.AddAudioTrack(mp4gomedia.MP4_CODEC_AAC, mp4gomedia.WithAudioChannelCount(1), mp4gomedia.WithAudioSampleRate(8000))

		/////////////////////////
		var inFilename = `..\container\res\output.aac`
		var fp io.ReadCloser
		if fp, err = os.Open(inFilename); err != nil {
			log.Println(err)
			return
		} else {
			defer fp.Close()
		}

		var (
			samples = 0
		)

		_ = formats.EmitADTSReader(fp, func(packet *mpeg4audio.ADTSPacket, data []byte, head []byte) {
			for isRun {
				if writeAudio.Load() {
					writeAudio.Store(false)
					break
				} else {
					time.Sleep(time.Millisecond)
				}
			}

			timestamp := time.Duration(samples*1000/packet.SampleRate) * time.Millisecond
			log.Println("audio", timestamp)

			pts := uint64(timestamp.Milliseconds())
			samples += 1024
			if err = m.Write(trackIdAudio, append(head, data...), pts, pts); err != nil {
				log.Println(err)
			}

			topPTS = pts
			writeVideo.Store(true)
		})
	}
	isRun = false
}

func testMP45() {
	f, _ := os.Create("out5.mp4")
	defer f.Close()

	m, err := formats.MP4GoMediaMuxer(f)
	if err != nil {
		log.Println(err)
		return
	} else {
		defer func() {
			if err := m.WriteTrailer(); err != nil {
				log.Println(err)
			} else {
				log.Println("close success")
			}
		}()
	}

	trackId := m.AddVideoTrack(mp4gomedia.MP4_CODEC_H264)

	inFilename := `..\container\res\input.h264`
	fp, err := os.Open(inFilename)
	if err != nil {
		log.Println(err)
		return
	} else {
		defer fp.Close()
	}

	var (
		timestamp time.Duration
		startCode = []byte{0x0, 0x0, 0x0, 0x01}
	)

	trackIdAudio := m.AddAudioTrack(mp4gomedia.MP4_CODEC_G711A, mp4gomedia.WithAudioChannelCount(1), mp4gomedia.WithAudioSampleRate(8000))

	inFilenameAudio := `..\container\res\output.g711a`
	fpAudio, err := os.Open(inFilenameAudio)
	if err != nil {
		log.Println(err)
		return
	} else {
		defer fpAudio.Close()
	}

	codecs.EmitNALUH264Reader(fp, codecs.NALUFormatAnnexB, codecs.NALUFormatNo, func(t h264.NALUType, data []byte) {
		pts := uint64(timestamp.Milliseconds())
		log.Println("video", timestamp)

		switch t {
		case h264.NALUTypeSPS, h264.NALUTypePPS:
			// 7
			buff := append(startCode, data...)

			if err = m.Write(trackId, buff, pts, pts); err != nil {
				log.Println(err)
				break
			}
		case h264.NALUTypeIDR, h264.NALUTypeNonIDR:
			// 1
			buff := append(startCode, data...)

			if err = m.Write(trackId, buff, pts, pts); err != nil {
				log.Println(err)
				break
			}

			{

				buff := make([]byte, 320)
				n, err := fpAudio.Read(buff)
				if err != nil {
					break
				}

				if err = m.Write(trackIdAudio, buff[:n], pts, pts); err != nil {
					log.Println(err)
					break
				}
			}

			timestamp += 40 * time.Millisecond
		default:
		}
	})

}

func testH264() {
	var (
		media *matroska.Matroska
		err   error
	)

	w, _ := os.Create("testh264.mkv")
	if media, err = matroska.Open(w, matroska.NewTrackH264()); err != nil {
		log.Println(err)
		return
	} else {
		defer media.Close()
	}

	stream := media

	/////////////////////////
	inFilename := `..\container\res\input.h264`
	var fp io.ReadCloser
	if fp, err = os.Open(inFilename); err != nil {
		log.Println(err)
		return
	} else {
		defer fp.Close()
	}

	var (
		timestamp time.Duration
		startCode = []byte{0x0, 0x0, 0x0, 0x01}
	)

	codecs.EmitNALUH264Reader(fp, codecs.NALUFormatAnnexB, codecs.NALUFormatNo, func(t h264.NALUType, data []byte) {
		switch t {
		case h264.NALUTypeSPS, h264.NALUTypePPS:
			// 7
			buff := append(startCode, data...)
			_, _ = stream.WriteVideo(timestamp, buff)
		case h264.NALUTypeIDR, h264.NALUTypeNonIDR:
			// 1
			buff := append(startCode, data...)
			_, _ = stream.WriteVideo(timestamp, buff)
			timestamp += 40 * time.Millisecond
		default:
		}
	})
}

func testHevc() {
	var (
		media *matroska.Matroska
		err   error
	)

	w, _ := os.Create("testHevc.mkv")
	if media, err = matroska.Open(w, matroska.NewTrackH265()); err != nil {
		log.Println(err)
		return
	} else {
		defer media.Close()
	}

	stream := media

	/////////////////////////
	inFilename := `..\container\res\480x272.h265`
	var fp io.ReadCloser
	if fp, err = os.Open(inFilename); err != nil {
		log.Println(err)
		return
	} else {
		defer fp.Close()
	}

	var (
		timestamp time.Duration
		startCode = []byte{0x0, 0x0, 0x0, 0x01}
	)

	codecs.EmitNALUH265Reader(fp, codecs.NALUFormatAnnexB, codecs.NALUFormatNo, func(t h265.NALUType, data []byte) {
		switch t {
		case h265.NALUType_VPS_NUT, h265.NALUType_SPS_NUT, h265.NALUType_PPS_NUT:
			buff := append(startCode, data...)
			_, _ = stream.WriteVideo(timestamp, buff)
		case h265.NALUType_IDR_W_RADL, h265.NALUType_TRAIL_R:
			buff := append(startCode, data...)
			_, _ = stream.WriteVideo(timestamp, buff)
			timestamp += 40 * time.Millisecond
		}
	})
}

func testAAC() {
	var (
		media *matroska.Matroska
		err   error
	)

	w, _ := os.Create("testAAC.mkv")
	if media, err = matroska.Open(w, matroska.NewTrackAAC(8000, 1)); err != nil {
		log.Println(err)
		return
	} else {
		defer media.Close()
	}

	stream := media

	/////////////////////////
	var inFilename = `..\container\res\output.aac`
	var fp io.ReadCloser
	if fp, err = os.Open(inFilename); err != nil {
		log.Println(err)
		return
	} else {
		defer fp.Close()
	}

	var (
		timestamp time.Duration
		samples   = 0
	)

	_ = formats.EmitADTSReader(fp, func(packet *mpeg4audio.ADTSPacket, data, _ []byte) {
		_, _ = stream.WriteAudio(timestamp, data)

		samples += 1024
		timestamp += time.Duration(samples*1000/packet.SampleRate) * time.Millisecond
	})
}

func testAAC2() {
	var (
		err        error
		inFilename = `..\container\res\output.aac`
		fp         io.ReadCloser

		media *matroska.Matroska
		track matroska.Track
	)

	if fp, err = os.Open(inFilename); err != nil {
		log.Println(err)
		return
	} else {
		defer fp.Close()
	}

	demuxer := formats.AACDemuxer(fp)
	{
		streams, err := demuxer.Streams()
		if err != nil {
			log.Println(err)
			return
		}

		for _, stream := range streams {
			if stream.Type() == av.AAC {
				switch stream := stream.(type) {
				case av.AudioCodecData:
					sampleRate := stream.SampleRate()
					channelLayout := int(stream.ChannelLayout())

					track = matroska.NewTrackAAC(sampleRate, channelLayout)
					break
				}
			}
		}
	}

	if track == nil {
		log.Println("not aac")
		return
	}

	w, _ := os.Create("testAAC2.mkv")
	if media, err = matroska.Open(w, track); err != nil {
		log.Println(err)
		return
	} else {
		defer media.Close()
	}

	for {
		packet, err := demuxer.ReadPacket()
		if err != nil {
			break
		}

		_, _ = media.WriteVideo(packet.Time, packet.Data)
	}
}

func testG711a() {
	var (
		media *matroska.Matroska
		err   error
	)

	w, _ := os.Create("testG711a.mkv")
	if media, err = matroska.Open(w, matroska.NewTrackPCMA(8000, 1)); err != nil {
		log.Println(err)
		return
	} else {
		defer media.Close()
	}

	stream := media

	/////////////////////////
	inFilename := `..\container\res\output.g711a`
	fp, err := os.Open(inFilename)
	if err != nil {
		log.Println(err)
		return
	} else {
		defer fp.Close()
	}

	var timestamp time.Duration

	for {
		buff := make([]byte, 320)
		n, err := fp.Read(buff)
		if err != nil {
			break
		}

		_, _ = stream.WriteAudio(timestamp, buff[:n])
		timestamp += 40 * time.Millisecond
	}
}

func testG726() {
	var (
		media *matroska.Matroska
		err   error
	)

	w, _ := os.Create("testG726.mkv")
	if media, err = matroska.Open(w, matroska.NewTrackG726(8000, 1)); err != nil {
		log.Println(err)
		return
	} else {
		defer media.Close()
	}

	stream := media

	/////////////////////////
	inFilename := `..\container\res\output.g726`
	var fp io.ReadCloser
	if fp, err = os.Open(inFilename); err != nil {
		log.Println(err)
		return
	} else {
		defer fp.Close()
	}

	var timestamp time.Duration

	for {
		buff := make([]byte, 160)
		n, err := fp.Read(buff)
		if err != nil {
			break
		}

		_, _ = stream.WriteAudio(timestamp, buff[:n])
		timestamp += 40 * time.Millisecond
	}
}

func testH264G711A() {
	var (
		media *matroska.Matroska
		err   error
	)

	w, _ := os.Create("testH264G711A.mkv")
	if media, err = matroska.Open(w, matroska.NewTrackH264(), matroska.NewTrackPCMA(8000, 1)); err != nil {
		log.Println(err)
		return
	} else {
		defer media.Close()
	}

	inAudioFilename := `..\container\res\output.g711a`
	var fpAudio io.ReadCloser
	if fpAudio, err = os.Open(inAudioFilename); err != nil {
		log.Println(err)
		return
	} else {
		defer fpAudio.Close()
	}

	inFilename := `..\container\res\input.h264`
	var fp io.ReadCloser
	if fp, err = os.Open(inFilename); err != nil {
		log.Println(err)
		return
	} else {
		defer fp.Close()
	}

	var (
		timestamp time.Duration
		startCode = []byte{0x0, 0x0, 0x0, 0x01}
	)

	codecs.EmitNALUH264Reader(fp, codecs.NALUFormatAnnexB, codecs.NALUFormatNo, func(t h264.NALUType, data []byte) {
		switch t {
		case h264.NALUTypeSPS, h264.NALUTypePPS:
			// 7
			buff := append(startCode, data...)
			_, _ = media.WriteVideo(timestamp, buff)
		case h264.NALUTypeIDR, h264.NALUTypeNonIDR:
			// 1
			{
				buff := make([]byte, 320)
				if n, err := fpAudio.Read(buff); err != nil {
					log.Println(err)
				} else {
					_, _ = media.WriteAudio(timestamp, buff[:n])
				}
			}

			{
				buff := append(startCode, data...)
				_, _ = media.WriteVideo(timestamp, buff)
			}

			timestamp += 40 * time.Millisecond

		default:
		}
	})
}

func rtspClient() {
	// github.com\bluenviron\gortsplib\v3@v3.10.0\examples\server
	// rtsp server

	var covertAVCCToAnnexB = func(data []byte) []byte {
		r := bytes.NewBuffer(data)

		out := bytes.Buffer{}
		_ = codecs.ConvertAVCCToAnnexB(r, &out)
		return out.Bytes()
	}

	filename := "input.mp4"
	fp, err := os.Open(filename)
	if err != nil {
		return
	} else {
		defer fp.Close()
	}

	demuxer := mp4vdk.NewDemuxer(fp)
	streams, err := demuxer.Streams()
	if err != nil {
		return
	}

	var (
		tracks    []*media.Media
		aacConfig aacparser.MPEG4AudioConfig
	)
	for _, stream := range streams {
		switch stream := stream.(type) {
		case h264parser.CodecData:
			tracks = append(tracks, &media.Media{
				Type: media.TypeVideo,
				Formats: []formatsvdk.Format{
					&formatsvdk.H264{
						PayloadTyp:        96,
						SPS:               stream.SPS(),
						PPS:               stream.PPS(),
						PacketizationMode: 1,
					},
				},
			})
		case h265parser.CodecData:
			tracks = append(tracks, &media.Media{
				Type: media.TypeVideo,
				Formats: []formatsvdk.Format{
					&formatsvdk.H265{
						PayloadTyp: 96,
						VPS:        stream.VPS(),
						SPS:        stream.SPS(),
						PPS:        stream.PPS(),
						MaxDONDiff: 0,
					},
				},
			})

		case aacparser.CodecData:
			aacConfig = stream.Config
			if true {
				tracks = append(tracks, &media.Media{
					Type: media.TypeAudio,
					Formats: []formatsvdk.Format{
						&formatsvdk.MPEG4Audio{
							PayloadTyp:     97,
							ProfileLevelID: 1,
							Config: &mpeg4audio.Config{
								Type:         mpeg4audio.ObjectTypeAACLC,
								SampleRate:   aacConfig.SampleRate,
								ChannelCount: aacConfig.ChannelLayout.Count(),
							},
							SizeLength:       13,
							IndexLength:      3,
							IndexDeltaLength: 3,
						},
					},
				})
			}
		}
	}

	var c = &gortsplib.Client{}

	rtspURL := fmt.Sprintf("rtsp://127.0.0.1:%v/%v", 8554, "aaa")

	err = c.StartRecording(rtspURL, tracks)
	if err != nil {
		log.Println(err)
		return
	}

	var (
		h264SSRC = uint32(222)
		aacSSRC  = uint32(223)
	)

	var h264Encoder = rtph264.Encoder{
		PayloadType: 96,
		SSRC:        &h264SSRC,
	}

	var aacEncoder = rtpmpeg4audiogeneric.Encoder{
		PayloadType:      97,
		SampleRate:       aacConfig.SampleRate,
		SSRC:             &aacSSRC,
		SizeLength:       13,
		IndexLength:      3,
		IndexDeltaLength: 3,
	}

	_ = h264Encoder.Init()
	_ = aacEncoder.Init()

	firstTime := time.Now()
	for {
		// 读取数据包
		pkt, err := demuxer.ReadPacket()
		if err != nil {
			log.Println(err)
			break
		}

		{
			duration := pkt.Time - time.Now().Sub(firstTime)
			log.Println(duration)
			if duration > 0 {
				time.Sleep(duration)
			}
		}

		if streams[pkt.Idx].Type().IsVideo() {
			// 转换格式
			data := covertAVCCToAnnexB(pkt.Data)

			// rtp打包
			packets, _ := h264Encoder.Encode([][]byte{data}, pkt.Time)
			for _, packet := range packets {
				if err = c.WritePacketRTP(tracks[0], packet); err != nil {
					log.Println(err)
				}
			}
		} else if streams[pkt.Idx].Type().IsAudio() {
			// rtp打包
			packets, err := aacEncoder.Encode([][]byte{pkt.Data}, pkt.Time)
			for _, packet := range packets {
				if err = c.WritePacketRTP(tracks[1], packet); err != nil {
					log.Println(err)
				}
			}
		}
	}

}

func ebmlDemo() {
	type VoidEbml struct {
		Void []byte `ebml:"Void"`
	}

	t := &VoidEbml{
		Void: make([]byte, 127),
	}

	buff := bytes.Buffer{}
	ebml.Marshal(t, &buff)
	log.Println(t)
}
